package ClassicalCryptography;

public class RailfenceCipher {
    private int row;
    private int column;

    public RailfenceCipher(int row) {
        this.row = row;
    }

    public String encrypt(String plaintext) {
        int len = plaintext.length();
        StringBuilder tempPlaintext = new StringBuilder(plaintext);
        if (len % row != 0) {
            for (int i = 0; i < row - len % row; i++) {
                tempPlaintext.append("*");
            }
        }
        column = tempPlaintext.length() / row;
        char[][] temp = new char[row][column];

        int charIndex = 0;
        for (int j = 0; j < column; j++) {
            for (int i = 0; i < row; i++) {
                temp[i][j] = tempPlaintext.charAt(charIndex++);
            }
        }
        StringBuilder res = new StringBuilder(tempPlaintext.length());
        for (int i = 0; i < row; i++) {
            res.append(temp[i]);
        }

        return res.substring(0);
    }

    public String decrypt(String ciphertext) {
        char[][] temp = new char[row][column];
        int charIndex = 0;
        for (int i = 0; i < row; i++) {
            for (int j = 0; j < column; j++) {
                temp[i][j] = ciphertext.charAt(charIndex++);
            }
        }

        StringBuilder res = new StringBuilder(charIndex);
        for (int j = 0; j < column; j++) {
            for (int i = 0; i < row; i++) {
                if (temp[i][j] != '*') {
                    res.append(temp[i][j]);
                }
            }
        }

        return res.substring(0);
    }
}
